package com.sxd.bill;

import java.util.Arrays;

import com.sxd.bill.common.item.MessageItem;
import com.sxd.bill.common.item.TextLabelItem;
import com.sxd.bill.ner.CRFNerWrapper;
import com.sxd.bill.ner.utils.OutputUtility;
import com.sxd.bill.rela.RelationExtractorManager;
import com.sxd.bill.rela.item.RelationContextItem;

public class NERRE {

    private CRFNerWrapper ner;
    private RelationExtractorManager relaManager;

    public NERRE(String modelPath) throws Exception {
        this.ner = new CRFNerWrapper(modelPath);
        this.relaManager = RelationExtractorManager.ofDefault();
    }

    public RelationContextItem extract(String text) {
        // 命名实体识别
        // TextLabelItem textLabelItem = ner.recognizeLabelEntities(text);
        TextLabelItem textLabelItem = ner.recognizeLabelEntities(text);
        // 关系抽取
        RelationContextItem relaContextItem = relaManager.extract_relation(textLabelItem);
        // 输出结果格式转换 TODO
        return relaContextItem;
    }

    public RelationContextItem extract(MessageItem msg) {
        String text = msg.getContent();
        // 命名实体识别
        // TextLabelItem textLabelItem = ner.recognizeLabelEntities(text);
        String[] tags = ner.recognize(text); // 得出模型预测每个字符发标签
        TextLabelItem textLabelItem = OutputUtility.getTextLabelItem(Arrays.asList(tags), text); // 根据字符标签获得要素实体
        // 关系抽取
        RelationContextItem relaContextItem = relaManager.extract_relation(textLabelItem);
        // 补充其他信息
        relaContextItem.setPublishDate(msg.getPubDate());
        relaContextItem.setPublisherNickName(msg.getPublisher());
        relaContextItem.setRoomName(msg.getRoomName());
        relaContextItem.setCrawlTime(msg.getCrawTime());
        return relaContextItem;
    }

}
